{
 "cells": [
  {
   "cell_type": "markdown",
   "metadata": {
    "id": "rsOnp1Y-TJy_"
   },
   "source": [
    "# Minimum Wage Example Notebook with DiD"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {
    "id": "trvqH1pjTJpR"
   },
   "source": [
    "This notebook implements Difference-in-Differences in an application on\n",
    "the effect of minimum wage changes on teen employment. We use data from\n",
    "[Callaway\n",
    "(2022)](https://bcallaway11.github.io/files/Callaway-Chapter-2022/main.pdf). The data are annual county level data from the United States covering 2001 to 2007. The outcome variable is log county-level teen employment, and the treatment variable is an indicator for whether the county has a minimum wage above the federal minimum wage. Note that this definition of the treatment variable makes the analysis straightforward but ignores the nuances of the exact value of the minimum wage in each county and how far those values are from the federal minimum. The data also include county population and county average annual pay.\n",
    "See [Callaway and Sant’Anna\n",
    "(2021)](https://www.sciencedirect.com/science/article/abs/pii/S0304407620303948)\n",
    "for additional details on the data.\n",
    "\n",
    "First, we will load some libraries."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {
    "colab": {
     "base_uri": "https://localhost:8080/"
    },
    "id": "KIcmAX3IGG-6",
    "outputId": "5b8ec3fe-3f86-42bc-80f5-cfa8477395d2"
   },
   "outputs": [],
   "source": [
    "!pip install doubleml~=0.7.0"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {
    "colab": {
     "base_uri": "https://localhost:8080/"
    },
    "id": "MER7xs5U0DRL",
    "outputId": "7694f2a3-9102-4999-f0cd-fca16d41c057"
   },
   "outputs": [],
   "source": [
    "!pip install formulaic"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {
    "id": "FFlG2QhXTJav"
   },
   "outputs": [],
   "source": [
    "import numpy as np\n",
    "import pandas as pd\n",
    "import doubleml as dml\n",
    "from sklearn.dummy import DummyRegressor, DummyClassifier\n",
    "from sklearn.linear_model import LinearRegression, LogisticRegression, LassoCV, RidgeCV, LogisticRegressionCV\n",
    "from sklearn.ensemble import RandomForestRegressor, RandomForestClassifier\n",
    "from sklearn.preprocessing import PolynomialFeatures\n",
    "from sklearn.model_selection import cross_val_predict, StratifiedKFold\n",
    "from sklearn.pipeline import make_pipeline\n",
    "import warnings\n",
    "\n",
    "warnings.filterwarnings(\"ignore\")\n",
    "np.random.seed(772023)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {
    "id": "TXQ_GfxZXDeG"
   },
   "outputs": [],
   "source": [
    "from sklearn.base import TransformerMixin, BaseEstimator\n",
    "from formulaic import Formula\n",
    "\n",
    "\n",
    "class FormulaTransformer(TransformerMixin, BaseEstimator):\n",
    "\n",
    "    def __init__(self, formula):\n",
    "        self.formula = formula\n",
    "\n",
    "    def fit(self, X, y=None):\n",
    "        return self\n",
    "\n",
    "    def transform(self, X, y=None):\n",
    "        df = Formula(self.formula).get_model_matrix(X)\n",
    "        return df.values.astype(np.float64)"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {
    "id": "u6jWjkrzU8I6"
   },
   "source": [
    "## Loading the data"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {
    "id": "znh8YcAXSp3E"
   },
   "outputs": [],
   "source": [
    "data = pd.read_csv(\"https://raw.githubusercontent.com/CausalAIBook/MetricsMLNotebooks/main/data/minwage_data.csv\",\n",
    "                   index_col=0)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {
    "colab": {
     "base_uri": "https://localhost:8080/",
     "height": 206
    },
    "id": "PQdsT6BnWKeq",
    "outputId": "2f92aa07-413b-4c00-9ed8-e143116b4942"
   },
   "outputs": [],
   "source": [
    "data.head()"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {
    "id": "v37g7zlwW5pH"
   },
   "source": [
    "### Data Preparation\n",
    "\n",
    "We remove observations that are already treated in the first observed period (2001). We drop all variables that we won't use in our analysis."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {
    "id": "W6ob7pptW49G"
   },
   "outputs": [],
   "source": [
    "data = data.loc[(data.G == 0) | (data.G > 2001)]\n",
    "data.drop(columns=[\"countyreal\", \"state_name\", \"FIPS\", \"emp0A01_BS\",\n",
    "                   \"quarter\", \"censusdiv\", \"pop\", \"annual_avg_pay\",\n",
    "                   \"state_mw\", \"fed_mw\", \"ever_treated\"], inplace=True)"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {
    "id": "Ri12EDNJaAfF"
   },
   "source": [
    "Next, we create the treatment groups. We focus our analysis exclusively on the set of counties that had wage increases away from the federal minimum wage in 2004. That is, we treat 2003 and earlier as the pre-treatment period."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {
    "id": "huj7huQ1aQSq"
   },
   "outputs": [],
   "source": [
    "years = [2001, 2002, 2003, 2004, 2005, 2006, 2007]\n",
    "treat, cont = {}, {}\n",
    "for year in years:\n",
    "    treat[year] = data.loc[(data.G == 2004) & (data.year == year)].copy()\n",
    "    cont[year] = data.loc[((data.G == 0) | (data.G > year)) & (data.year == year)].copy()"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {
    "id": "6HC1PX_Uc5bQ"
   },
   "source": [
    "We assume that the basic assumptions, particularly parallel trends, hold after conditioning on pre-treatment variables: 2001 population, 2001 average pay and 2001 teen employment, as well as the region in which the county is located. (The region is characterized by four\n",
    "categories.)\n",
    "\n",
    "Consequently, we want to extract the control variables for both treatment and control group in 2001."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {
    "id": "KvkwAdL6evsU"
   },
   "outputs": [],
   "source": [
    "treat[2001].drop(columns=[\"year\", \"G\", \"region\", \"treated\"], inplace=True)\n",
    "cont[2001].drop(columns=[\"year\", \"G\", \"region\", \"treated\"], inplace=True)"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {
    "id": "zU7rM_5Ne3Xr"
   },
   "source": [
    "2003 serves as the pre-treatment period for both counties that do receive the treatment in 2004 and those that do not."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {
    "id": "3cd3dBDqeyqa"
   },
   "outputs": [],
   "source": [
    "treatB = pd.merge(treat[2003], treat[2001], on=\"id\", suffixes=[\"_pre\", \"_0\"])\n",
    "treatB.drop(columns=[\"treated\", \"lpop_pre\", \"lavg_pay_pre\", \"year\", \"G\"], inplace=True)\n",
    "\n",
    "contB = pd.merge(cont[2003], cont[2001], on=\"id\", suffixes=[\"_pre\", \"_0\"])\n",
    "contB.drop(columns=[\"treated\", \"lpop_pre\", \"lavg_pay_pre\", \"year\", \"G\"], inplace=True)"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {
    "id": "xL1fSfb5e82d"
   },
   "source": [
    "We estimate the ATET in 2004-2007, which corresponds to the effect in the year of treatment as well as in the three years after the treatment. The control observations are the observations that still have the federal minimum wage in each year. (The control group is shrinking in each year as additional units receive treatment)."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {
    "id": "zvN6Nmy0gPy4"
   },
   "outputs": [],
   "source": [
    "tdid, cdid = {}, {}\n",
    "did_data = {}\n",
    "# the first year will be used for pre-trend testing later on\n",
    "for year in [2002, 2004, 2005, 2006, 2007]:\n",
    "    treat[year].drop(columns=[\"lpop\", \"lavg_pay\", \"year\", \"G\", \"region\"], inplace=True)\n",
    "    cont[year].drop(columns=[\"lpop\", \"lavg_pay\", \"year\", \"G\", \"region\"], inplace=True)\n",
    "\n",
    "    tdid[year] = pd.merge(treat[year], treatB, on=\"id\")\n",
    "    tdid[year][\"dy\"] = tdid[year][\"lemp\"] - tdid[year][\"lemp_pre\"]\n",
    "    tdid[year].drop(columns=[\"id\", \"lemp\", \"lemp_pre\"], inplace=True)\n",
    "    tdid[year].treated = 1  # forcing treatment to be 1, so that 2002 is \"treated\" when testing pre-trends\n",
    "\n",
    "    cdid[year] = pd.merge(cont[year], contB, on=\"id\")\n",
    "    cdid[year][\"dy\"] = cdid[year][\"lemp\"] - cdid[year][\"lemp_pre\"]\n",
    "    cdid[year].drop(columns=[\"id\", \"lemp\", \"lemp_pre\"], inplace=True)\n",
    "\n",
    "    # join control and treatment data and create one-hot-encoded region variables\n",
    "    did_data[year] = pd.concat((tdid[year], cdid[year]))\n",
    "    dummy_data = pd.get_dummies(did_data[year].region, drop_first=True, prefix=\"region\")\n",
    "    did_data[year] = pd.concat((did_data[year], dummy_data), axis=1).drop(columns=[\"region\"])"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {
    "colab": {
     "base_uri": "https://localhost:8080/",
     "height": 206
    },
    "id": "RhFOwd40RQ6K",
    "outputId": "32348085-54c1-48d2-8500-df15541331d7"
   },
   "outputs": [],
   "source": [
    "did_data[2004].head()"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {
    "id": "EqHmiHaZgPZz"
   },
   "source": [
    "### Estimation of the ATET with DML\n",
    "\n",
    "We estimate the ATET of the county level minimum wage being larger than the federal minimum with the DML algorithm presented in Section 16.3 in the book. This requires estimation of the nuisance functions $E[Y|D=0,X]$, $E[D|X]$ as well as $P(D = 1)$. For the conditional expectation functions, we will consider different modern ML regression methods, namely: Constant (= no controls); a linear combination of the controls; an expansion of the raw control variables including all third order interactions; Lasso (CV); Ridge (CV); Random Forest.\n",
    "The methods indicated with CV have their tuning parameter selected by cross-validation.\n"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {
    "id": "5WjxpBKUGG_C"
   },
   "source": [
    "The following code block implements the DML estimator with cross-fitting"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {
    "id": "Ut2Rrt-5GG_C"
   },
   "outputs": [],
   "source": [
    "def final_stage(D, y, Dhat, yhat, phat):\n",
    "    # doubly robust quantity for every sample\n",
    "    phihat = ((D - Dhat) / (phat * (1 - Dhat))) * (y - yhat)\n",
    "    point = np.mean(phihat) / np.mean(D / phat)\n",
    "    # influence function\n",
    "    phihat = (phihat - point * (D / phat)) / np.mean(D / phat)\n",
    "    var = np.mean(np.square(phihat))\n",
    "    stderr = np.sqrt(var / D.shape[0])\n",
    "    return point, stderr\n",
    "\n",
    "\n",
    "def dr_att(X, D, y, modely, modeld, *, trimming=0.01, nfolds=5):\n",
    "    '''\n",
    "    DML for the ATT estimation (Doubly Robust Learning)\n",
    "    with cross-fitting\n",
    "\n",
    "    Input\n",
    "    -----\n",
    "    X: the controls, pandas DataFrame\n",
    "    D: treatment indicator, numpy array\n",
    "    y: the outcome (the delta before and after intervention in DiD), numpy array\n",
    "    modely: the ML model for predicting y (the delta Y_1 - Y_0 in DiD)\n",
    "    modeld: the ML model for predicting the treatment D\n",
    "    trimming: threshold below which to trim propensities, float (default=0.01)\n",
    "    nfolds: the number of folds in cross-fitting, int (default=5)\n",
    "\n",
    "    Output\n",
    "    ------\n",
    "    point: the point estimate of the treatment on the treated\n",
    "    stderr: the standard error of the treatment effect\n",
    "    yhat: cross-fitted predictions of the outcome y under control, E[y|D=0, X]\n",
    "    Dhat: cross-fitted predictions of the treatment D, E[D|X]\n",
    "    rmsey: the RMSE of the model y ~ X | D==0\n",
    "    rmseD: the RMSE of the model D ~ X\n",
    "    phat: the estimated treatment probability in each fold\n",
    "    '''\n",
    "    cv = StratifiedKFold(n_splits=nfolds, shuffle=True, random_state=1234)\n",
    "    # fit a model E[y | D=0, X]\n",
    "    yhat = np.zeros(y.shape)\n",
    "    for train, test in cv.split(X, D):\n",
    "        modely.fit(X.iloc[train][D[train] == 0], y[train][D[train] == 0])\n",
    "        yhat[test] = modely.predict(X.iloc[test])\n",
    "    # fit a model P[D]\n",
    "    phat = cross_val_predict(DummyRegressor(), X, D, cv=cv)\n",
    "    # propensity scores E[D | X]\n",
    "    Dhat = cross_val_predict(modeld, X, D, cv=cv, method='predict_proba')[:, 1]\n",
    "    # trimm propensity score\n",
    "    Dhat = np.clip(Dhat, trimming, 1 - trimming)\n",
    "    # estimation of parameter\n",
    "    point, stderr = final_stage(D, y, Dhat, yhat, phat)\n",
    "    # nuisance function rmse's\n",
    "    rmsey = np.sqrt(np.mean((y - yhat)[D == 0]**2))\n",
    "    rmseD = np.sqrt(np.mean((D - Dhat)**2))\n",
    "    return point, stderr, yhat, Dhat, rmsey, rmseD, phat"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {
    "id": "6lBHCpuKXtu5"
   },
   "outputs": [],
   "source": [
    "def get_nuisance_learners(region_names):\n",
    "    ''' Constructs the learners we will consider for nuisance estimation\n",
    "\n",
    "    region_names: the available region names\n",
    "\n",
    "    Returns: a dictionary of learners of the form {key: (ml_g, ml_m)}\n",
    "    '''\n",
    "    # a formula that will be used within a featurizer to create interactions\n",
    "    # (region id) * (baseline outcome) + (region id) + (baseline outcomes)\n",
    "    formula = '0 + ' + ' + '.join([f\"{col} * (lemp_0 + lpop_0 + lavg_pay_0)\" for col in region_names])\n",
    "\n",
    "    # a generic polynomial featurizer\n",
    "    def poly():\n",
    "        return PolynomialFeatures(degree=3, interaction_only=True, include_bias=False)\n",
    "\n",
    "    learners = {\n",
    "        # no X\n",
    "        \"No Controls\": (DummyRegressor(strategy=\"mean\"), DummyClassifier(strategy=\"prior\")),\n",
    "        # linear models of X\n",
    "        \"Basic\": (LinearRegression(), LogisticRegression(random_state=123)),\n",
    "        # linear models with region interactions\n",
    "        \"Expansion\": (make_pipeline(FormulaTransformer(formula), LinearRegression()),\n",
    "                      make_pipeline(FormulaTransformer(formula), LogisticRegression(random_state=123))),\n",
    "        # full blown third degree interactions of X, with Lasso\n",
    "        \"Lasso (CV)\": (make_pipeline(poly(), LassoCV(n_jobs=-1, random_state=123)),\n",
    "                       make_pipeline(poly(), LogisticRegressionCV(penalty=\"l1\", solver=\"liblinear\",\n",
    "                                                                  n_jobs=-1, random_state=123))),\n",
    "        # full blown third degree interactions, with Ridge\n",
    "        \"Ridge (CV)\": (make_pipeline(poly(), RidgeCV()),\n",
    "                       make_pipeline(poly(), LogisticRegressionCV(n_jobs=-1, random_state=123))),\n",
    "        # Random Forest regressors and classifiers\n",
    "        \"Random Forest\": (RandomForestRegressor(n_estimators=100, min_samples_leaf=20, max_features=4,\n",
    "                                                n_jobs=-1, random_state=123),\n",
    "                          RandomForestClassifier(n_estimators=100, min_samples_leaf=20, max_features=4,\n",
    "                                                 n_jobs=-1, random_state=123))\n",
    "    }\n",
    "    return learners"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {
    "colab": {
     "base_uri": "https://localhost:8080/"
    },
    "id": "ZQd7HP57GG_D",
    "outputId": "d51e056b-e7c1-45b9-d6fb-0789e17fd8cb"
   },
   "outputs": [],
   "source": [
    "def MinWageDiD(years):\n",
    "    # arrays for saving the results\n",
    "    att, se_att, RMSE_d, RMSE_y = {}, {}, {}, {}\n",
    "\n",
    "    for year in years:\n",
    "        print(f\"Estimating ATET for year {year}. Please wait.\")\n",
    "        att[year], se_att[year], RMSE_d[year], RMSE_y[year] = {}, {}, {}, {}\n",
    "\n",
    "        # set up the data for the specific year\n",
    "        X = did_data[year].drop(columns=[\"treated\", \"dy\"])  # controls\n",
    "        D = did_data[year].treated.values  # treatment\n",
    "        dy = did_data[year].dy.values  # delta outcome betwee post and pre treatment\n",
    "        region_names = [col for col in X.columns if col.startswith('region_')]\n",
    "\n",
    "        # get dictionary of nuisance learners\n",
    "        learners = get_nuisance_learners(region_names)\n",
    "\n",
    "        # for storing the nuisance predictions\n",
    "        pred_y, pred_d = {}, {}\n",
    "        for method, (ml_g, ml_m) in learners.items():\n",
    "            print(f\"Estimating {method}.\")\n",
    "            point, stderr, dyhat, Dhat, rmsey, rmsed, phat = dr_att(X, D, dy, ml_g, ml_m, trimming=0.01, nfolds=5)\n",
    "            pred_y[method], pred_d[method] = dyhat, Dhat\n",
    "            att[year][method], se_att[year][method] = point, stderr\n",
    "            RMSE_y[year][method], RMSE_d[year][method] = rmsey, rmsed\n",
    "\n",
    "        # find best model for dy and d and use predictions\n",
    "        besty = min(RMSE_y[year], key=RMSE_y[year].get)\n",
    "        bestd = min(RMSE_d[year], key=RMSE_d[year].get)\n",
    "        # phat this is the same for all methods and equal to the cross-fitted Pr[D=1]\n",
    "        point, stderr = final_stage(D, dy, pred_d[bestd], pred_y[besty], phat)\n",
    "        att[year][\"Best\"], se_att[year][\"Best\"] = point, stderr\n",
    "        RMSE_y[year][\"Best\"] = RMSE_y[year][besty]\n",
    "        RMSE_d[year][\"Best\"] = RMSE_d[year][bestd]\n",
    "\n",
    "        # we can also find the best linear combination of models via stacking\n",
    "        pred_y = pd.DataFrame(pred_y)\n",
    "        pred_d = pd.DataFrame(pred_d)\n",
    "        stacked_pred_y = LinearRegression().fit(pred_y, dy).predict(pred_y)\n",
    "        stacked_pred_d = LinearRegression().fit(pred_d, D).predict(pred_d)\n",
    "        # phat this is the same for all methods and equal to the cross-fitted Pr[D=1]\n",
    "        point, stderr = final_stage(D, dy, stacked_pred_d, stacked_pred_y, phat)\n",
    "        att[year][\"Stack\"], se_att[year][\"Stack\"] = point, stderr\n",
    "        RMSE_y[year][\"Stack\"] = np.sqrt(np.mean((dy - stacked_pred_y)[D == 0]**2))\n",
    "        RMSE_d[year][\"Stack\"] = np.sqrt(np.mean((D - stacked_pred_d)**2))\n",
    "\n",
    "    return att, se_att, RMSE_y, RMSE_d\n",
    "\n",
    "\n",
    "att, se_att, RMSE_y, RMSE_d = MinWageDiD([2004, 2005, 2006, 2007])"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {
    "id": "TIwONJ0hGG_E"
   },
   "source": [
    "We start by reporting the RMSE obtained during cross-fitting for each learner in each period."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {
    "colab": {
     "base_uri": "https://localhost:8080/",
     "height": 332
    },
    "id": "592YvTt0GG_E",
    "outputId": "ca6b49a3-8177-4e39-a276-0e9f5d6815a3"
   },
   "outputs": [],
   "source": [
    "table1 = pd.concat({'RMSE dy': pd.DataFrame(RMSE_y),\n",
    "                    'RMSE D': pd.DataFrame(RMSE_d)}, axis=1)\n",
    "table1 = table1.swaplevel(0, 1, axis=1)\n",
    "table1 = table1.sort_index(axis=1)\n",
    "table1"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {
    "id": "Sea-bOaoGG_F"
   },
   "source": [
    "It appears there is some signal in the regressors, as all methods outside of  LassoCV produce somewhat smaller RMSEs than the No Controls baseline. While it would be hard to reliably conclude which of the relatively good performing methods is statistically best here, Best (or a different ensemble) provides a good baseline that is principled in the sense that one could pre-commit to using the best learners without having first looked at the subsequent estimation results."
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {
    "id": "ogTzW6JLGG_F"
   },
   "source": [
    "We report estimates of the ATET in each period in the following table, together with the aforementioned RMSE's."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {
    "colab": {
     "base_uri": "https://localhost:8080/",
     "height": 332
    },
    "id": "iPX5OnYdAb8c",
    "outputId": "809ea250-2011-4e21-b67d-9a7b1746ce56"
   },
   "outputs": [],
   "source": [
    "table2 = pd.concat({'att': pd.DataFrame(att),\n",
    "                    'se': pd.DataFrame(se_att),\n",
    "                    'RMSE dy': pd.DataFrame(RMSE_y),\n",
    "                    'RMSE D': pd.DataFrame(RMSE_d)}, axis=1)\n",
    "table2 = table2.swaplevel(0, 1, axis=1)\n",
    "table2 = table2.sort_index(axis=1)\n",
    "table2"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {
    "id": "b42a0MLLGG_G"
   },
   "source": [
    "Here, we see that most methods provide point estimates that suggest the effect of the minimum wage increase leads to decreases in youth employment with small effects in the initial period that become larger in the years following the treatment. This pattern seems economically plausible as it may take time for firms to adjust employment and other input choices in response to a minimum wage change. In the estimates that are reported in the book we have values that are not consistent with this pattern, however, they systematically underperform in terms of having poor cross-fit prediction performance. In terms of point estimates, the other pattern that emerges is that all estimates that use the covariates produce ATET estimates that are systematically smaller in magnitude than the No Controls baseline, suggesting that failing to include the controls may lead to overstatement of treatment effects in this example.\n",
    "\n",
    "Turning to inference, we would reject the hypothesis of no minimum wage effect two or more years after the change at the 5% level, even after multiple testing correction, if we were to focus on many of the estimators."
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {
    "id": "OKKUIT8SGG_G"
   },
   "source": [
    "### Assess pre-trends\n",
    "\n",
    "Because we have data for the period 2001-2007, we can perform a so-called pre-trends test to provide some evidence about the plausibility of the conditional parallel trends assumption. Specifically, we can continue to use 2003 as the reference period but now consider 2002 to be the treatment period. Sensible economic mechanisms underlying the assumption would then typically suggest that the ATET in 2002 - before the 2004 minimum wage change we are considering - should be zero. Finding evidence that the ATET in 2002 is non-zero then calls into question the validity of the assumption.\n",
    "\n",
    "We change the treatment status of those observations, which received treatment in 2004 in the 2002 data and create a placebo treatment as well as control group. We call the same estimation pipeline that we used to obtain our ATET estimates for 2004-2007. We now hope that we get a statistical null result."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {
    "colab": {
     "base_uri": "https://localhost:8080/"
    },
    "id": "dOw9mzvEIBZ7",
    "outputId": "44c4c877-b6a7-4254-95f0-6d9726cce9fe"
   },
   "outputs": [],
   "source": [
    "att, se_att, RMSE_y, RMSE_d = MinWageDiD([2002])"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {
    "colab": {
     "base_uri": "https://localhost:8080/",
     "height": 332
    },
    "id": "NXBR0GnVJEKz",
    "outputId": "07965303-f44a-4923-972a-a0e54924fb7e"
   },
   "outputs": [],
   "source": [
    "table2002 = pd.concat({'att': pd.DataFrame(att),\n",
    "                       'se': pd.DataFrame(se_att),\n",
    "                       'RMSE dy': pd.DataFrame(RMSE_y),\n",
    "                       'RMSE D': pd.DataFrame(RMSE_d)}, axis=1)\n",
    "table2002 = table2002.swaplevel(0, 1, axis=1)\n",
    "table2002 = table2002.sort_index(axis=1)\n",
    "table2002"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {
    "id": "up95SCKAGG_J"
   },
   "source": [
    "Here we see broad agreement across all methods in the sense of returning point estimates that are small in magnitude and small relative to standard errors. In no case would we reject the hypothesis that the pre-event effect in 2002 is different from zero at usual levels of significance. We note that failing to reject the hypothesis of no pre-event effects certainly does not imply that the conditional DiD assumption is in fact satisfied. For example, confidence intervals include values that would be consistent with relatively large pre-event effects. However, it is reassuring to see that there is not strong evidence of a violation of the underlying identifying assumption."
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {
    "id": "vj0RZ1uu5EDO"
   },
   "source": [
    "### Visualizing the Effect of Treatment in 2004 Over the Years"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {
    "colab": {
     "base_uri": "https://localhost:8080/",
     "height": 472
    },
    "id": "ZGoHLiTa5C3H",
    "outputId": "30b364cb-915d-450a-faa1-a795006196d6"
   },
   "outputs": [],
   "source": [
    "import matplotlib.pyplot as plt\n",
    "\n",
    "table_all = pd.concat([table2002, table2], axis=1)\n",
    "years = np.array([2002, 2004, 2005, 2006, 2007])\n",
    "for it, method in enumerate(['No Controls', 'Best', 'Stack']):\n",
    "    plt.errorbar(years + (it - 1) * .1,\n",
    "                 np.array([table_all[year]['att'].loc[method] for year in years]),\n",
    "                 yerr=1.96 * np.array([table_all[year]['se'].loc[method] for year in years]),\n",
    "                 fmt='o', label=method)\n",
    "plt.legend()\n",
    "plt.xlabel('Year')\n",
    "plt.ylabel('ATET')\n",
    "plt.title('Effect of Treatment in 2004 Over the Years')\n",
    "plt.show()"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {
    "id": "wqh57pA0GG_J"
   },
   "source": [
    "### Repication with `DoubleML`\n",
    "The high-level implementation `DoubleML` can be used to replicate the above. Please note, that it will run for a while (around 15min)."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {
    "colab": {
     "background_save": true
    },
    "id": "jX6O2JB4PB2i"
   },
   "outputs": [],
   "source": [
    "def MinWageDiDwithDoubleML(years):\n",
    "    # arrays for saving the results\n",
    "    att, se_att, RMSE_d, RMSE_y = {}, {}, {}, {}\n",
    "\n",
    "    for year in years:\n",
    "        print(f\"Estimating ATET for year {year}. Please wait.\")\n",
    "        att[year], se_att[year], RMSE_d[year], RMSE_y[year] = {}, {}, {}, {}\n",
    "\n",
    "        # set up the data for the specific year\n",
    "        region_names = [col for col in did_data[year].columns if col.startswith(\"region_\")]\n",
    "        dml_data = dml.DoubleMLData(data=did_data[year],\n",
    "                                    x_cols=[\"lemp_0\", \"lpop_0\", \"lavg_pay_0\"] + region_names,\n",
    "                                    y_col=\"dy\",\n",
    "                                    d_cols=\"treated\")\n",
    "\n",
    "        # get dictionary of nuisance learners\n",
    "        learners = get_nuisance_learners(region_names)\n",
    "\n",
    "        for method, (ml_g, ml_m) in learners.items():\n",
    "            print(f\"Estimating {method}.\")\n",
    "            if method == \"Expansion\":\n",
    "                # we can't use the FormulaTransformer pipeline with DoubleML because\n",
    "                # it strips internally the pandas metadata\n",
    "                # so we have to explicitly transform the data before passing to DoubleML\n",
    "                trans_dml_data = dml.DoubleMLData.from_arrays(x=ml_g[0].fit_transform(did_data[year]),\n",
    "                                                              y=did_data[year].dy.values,\n",
    "                                                              d=did_data[year].treated.values)\n",
    "                dml_obj = dml.DoubleMLDID(trans_dml_data, ml_g=ml_g[1], ml_m=ml_m[1],\n",
    "                                          in_sample_normalization=False, trimming_threshold=0.01)\n",
    "            else:\n",
    "                dml_obj = dml.DoubleMLDID(dml_data, ml_g=ml_g, ml_m=ml_m,\n",
    "                                          in_sample_normalization=False, trimming_threshold=0.01)\n",
    "            # The following two lines are only needed to exactly replicate the custom\n",
    "            # implementation results, so that the same k-folds are used in cross-fitting\n",
    "            cv = StratifiedKFold(n_splits=5, shuffle=True, random_state=1234)\n",
    "            dml_obj.set_sample_splitting(list(cv.split(did_data[year][dml_data.x_cols],\n",
    "                                                       did_data[year].treated.values)))\n",
    "            dml_obj.fit()\n",
    "            att[year][method] = dml_obj._coef[0]\n",
    "            se_att[year][method] = dml_obj._se[0]\n",
    "            RMSE_y[year][method] = np.mean(dml_obj.rmses[\"ml_g0\"])\n",
    "            RMSE_d[year][method] = np.mean(dml_obj.rmses[\"ml_m\"])\n",
    "\n",
    "    return att, se_att, RMSE_y, RMSE_d"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {
    "colab": {
     "background_save": true
    },
    "id": "uVBmKdMkbSmW",
    "outputId": "d8b75835-aed7-4033-931f-f5d46da92ebb"
   },
   "outputs": [],
   "source": [
    "att, se_att, RMSE_y, RMSE_d = MinWageDiDwithDoubleML([2004, 2005])"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {
    "colab": {
     "background_save": true
    },
    "id": "6blwjZM9Xt2f",
    "outputId": "48d9c39d-4cae-4213-f613-c22d23763d07"
   },
   "outputs": [],
   "source": [
    "tabledml = pd.concat({'att': pd.DataFrame(att),\n",
    "                      'se': pd.DataFrame(se_att),\n",
    "                      'RMSE dy': pd.DataFrame(RMSE_y),\n",
    "                      'RMSE D': pd.DataFrame(RMSE_d)}, axis=1)\n",
    "tabledml = tabledml.swaplevel(0, 1, axis=1)\n",
    "tabledml = tabledml.sort_index(axis=1)\n",
    "tabledml"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {
    "id": "A1BRGg3neyka"
   },
   "source": [
    "We find that the DoubleML package returns basically the same results as the custom implementation (presented below for comparison), modulo the small randomness due to the random sample splitting involved in the cross-fitting process."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {
    "colab": {
     "background_save": true
    },
    "id": "bN1pGRFIe3Gc",
    "outputId": "61c83b15-5533-44e4-c945-16ae9837a629"
   },
   "outputs": [],
   "source": [
    "table2"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {
    "id": "SgUakV3TbbH8"
   },
   "source": []
  }
 ],
 "metadata": {
  "accelerator": "GPU",
  "colab": {
   "gpuType": "A100",
   "machine_shape": "hm",
   "provenance": []
  },
  "kernelspec": {
   "display_name": "Python 3 (ipykernel)",
   "language": "python",
   "name": "python3"
  },
  "language_info": {
   "codemirror_mode": {
    "name": "ipython",
    "version": 3
   },
   "file_extension": ".py",
   "mimetype": "text/x-python",
   "name": "python",
   "nbconvert_exporter": "python",
   "pygments_lexer": "ipython3",
   "version": "3.11.5"
  }
 },
 "nbformat": 4,
 "nbformat_minor": 1
}
